// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <pollux/dwio/common/reader.h>
#include <pollux/dwio/common/reader_factory.h>
#include <pollux/dwio/parquet/reader/metadata.h>
#include <pollux/dwio/parquet/reader/parquet_stats_context.h>

namespace kumo::pollux::dwio::common {
    class SelectiveColumnReader;
    class BufferedInput;
} // namespace kumo::pollux::dwio::common

namespace kumo::pollux::parquet {
    enum class ParquetMetricsType { HEADER, FILE_METADATA, FILE, BLOCK, TEST };

    class StructColumnReader;

    class ReaderBase;

    /// Implements the RowReader interface for Parquet.
    class ParquetRowReader : public dwio::common::RowReader {
    public:
        ParquetRowReader(
            const std::shared_ptr<ReaderBase> &readerBase,
            const dwio::common::RowReaderOptions &options);

        ~ParquetRowReader() override = default;

        int64_t nextRowNumber() override;

        int64_t nextReadSize(uint64_t size) override;

        uint64_t next(
            uint64_t size,
            pollux::VectorPtr &result,
            const dwio::common::Mutation * = nullptr) override;

        void updateRuntimeStats(
            dwio::common::RuntimeStatistics &stats) const override;

        void resetFilterCaches() override;

        std::optional<size_t> estimatedRowSize() const override;

        bool allPrefetchIssued() const override {
            //  Allow opening the next split while this is reading.
            return true;
        }

        // Checks if the specific row group is buffered.
        // Returns false if the row group is not loaded into buffer
        // or the buffered data has been evicted.
        bool isRowGroupBuffered(int32_t rowGroupIndex) const;

    private:
        // Compares row group  metadata to filters in ScanSpec in options of
        // ReaderBase and determines the set of row groups to scan.
        void filterRowGroups();

    protected:
        class Impl;
        std::unique_ptr<Impl> impl_;
    };

    /// Implements the reader interface for Parquet.
    class ParquetReader : public dwio::common::Reader {
    public:
        ParquetReader(
            std::unique_ptr<dwio::common::BufferedInput>,
            const dwio::common::ReaderOptions &options);

        ~ParquetReader() override = default;

        std::optional<uint64_t> numberOfRows() const override;

        std::unique_ptr<dwio::common::ColumnStatistics> columnStatistics(
            uint32_t index) const override {
            return nullptr;
        }

        const pollux::RowTypePtr &rowType() const override;

        const std::shared_ptr<const dwio::common::TypeWithId> &typeWithId()
        const override;

        std::unique_ptr<dwio::common::RowReader> createRowReader(
            const dwio::common::RowReaderOptions &options = {}) const override;

        FileMetaDataPtr fileMetaData() const;

    private:
        std::shared_ptr<ReaderBase> readerBase_;
    };

    class ParquetReaderFactory : public dwio::common::ReaderFactory {
    public:
        ParquetReaderFactory() : ReaderFactory(dwio::common::FileFormat::PARQUET) {
        }

        std::unique_ptr<dwio::common::Reader> createReader(
            std::unique_ptr<dwio::common::BufferedInput> input,
            const dwio::common::ReaderOptions &options) override {
            return std::make_unique<ParquetReader>(std::move(input), options);
        }
    };

    void registerParquetReaderFactory();

    void unregisterParquetReaderFactory();
} // namespace kumo::pollux::parquet
